{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "9413e752-a21c-4bd6-823c-cca04843ace0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "63"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "GENOMES=[\n",
    "    \"A188\",\"A632\",\"B97\",\"B104\",\"CIMBL55\",\"CML52\",\"CML69\",\"CML103\",\"CML228\",\n",
    "    \"CML247\",\"CML277\",\"CML322\",\"CML333\",\"CML457\",\"CML459\",\"CML530\",\"Chang7-2\",\n",
    "    \"DK105\",\"Dan340\",\"EP1\",\"F7\",\"HP301\",\"Huangzaosi\",\"Ia453\",\"Il14H\",\"Jing92\",\n",
    "    \"Jing724\",\"K0326Y\",\"Ki3\",\"Ki11\",\"Ky21\",\"LH244\",\"M37W\",\"M162W\",\"Mo17\",\"Mo18W\",\n",
    "    \"Ms71\",\"NC350\",\"NC358\",\"Oh7B\",\"Oh43\",\"P39\",\"PDJ\",\"PE0075\",\"S37\",\"Tx303\",\n",
    "    \"SK\",\"Tzi8\",\"Xu178\",\"Ye478\",\"Zheng58\",\"PI615697\",\"Ames21814\",\"TIL01\",\"TIL11\",\n",
    "    \"PI566673\",\"TIL18\",\"TIL25\",\"PT\",\"TAB\",\"ZAP\",\"PH207\",\"W22\",\n",
    "]\n",
    "\n",
    "len(GENOMES)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "9166725a-4b14-46a5-b7fd-2f5955646aae",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "sv_id = \"A188-1-797667\"\n",
    "block_file = \"./files/A188-1-797667/haploview.GABRIELblocks\"\n",
    "info_file = \"./files/A188-1-797667/SV.plink.info\"\n",
    "tag_file = \"./files/A188-1-797667/haploview.TAGS\"\n",
    "\n",
    "all_snp_df = pd.DataFrame()\n",
    "all_block_context = []\n",
    "\n",
    "if os.path.exists(block_file) and os.path.exists(info_file):\n",
    "    info_df = pd.read_csv(info_file, header=None, sep=\"\\t\")\n",
    "    block_count = 0\n",
    "    with open(block_file, \"r\") as f:\n",
    "        lines = f.readlines()\n",
    "        \n",
    "        block_context = []\n",
    "        block_id = \"\"\n",
    "        \n",
    "        for line in lines:\n",
    "            line_text = line.strip()\n",
    "            if \"BLOCK\" in line_text:\n",
    "                block_count+=1\n",
    "                block_id = f\"{sv_id}-{block_count}\"\n",
    "                line_list = line_text.split(\" \")\n",
    "                marker_index = line_list.index(\"MARKERS:\")\n",
    "                snp_indexs = [int(v)-1 for v in line_list[marker_index+1:]]\n",
    "                item_snp_df = info_df.iloc[snp_indexs][[0]].copy()\n",
    "                # item_snp_df[0] = item_snp_df[0].apply(lambda x: x.split(\"_\")[0])\n",
    "                item_snp_df[1] = block_id\n",
    "\n",
    "                if len(all_snp_df) > 0:\n",
    "                    all_snp_df = pd.concat([all_snp_df, item_snp_df])\n",
    "                else:\n",
    "                    all_snp_df = item_snp_df\n",
    "\n",
    "                if len(block_context) > 0:\n",
    "                    all_block_context += block_context\n",
    "                    block_context = []\n",
    "            else:\n",
    "                if len(block_context) == 0 and block_id !=\"\":\n",
    "                    block_context.append(f\">{block_id}\")\n",
    "                    block_context.append(line_text)\n",
    "                else:\n",
    "                    block_context.append(line_text)\n",
    "        if len(block_context) > 0:\n",
    "            all_block_context += block_context\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "658ac080-ab05-4a9f-bc22-e5a14135dde9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>0</th>\n",
       "      <th>1</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>34</th>\n",
       "      <td>1_147027</td>\n",
       "      <td>A188-1-797667-1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>36</th>\n",
       "      <td>1_147057</td>\n",
       "      <td>A188-1-797667-1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>37</th>\n",
       "      <td>1_147071</td>\n",
       "      <td>A188-1-797667-1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>39</th>\n",
       "      <td>1_147088</td>\n",
       "      <td>A188-1-797667-1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>46</th>\n",
       "      <td>1_148314</td>\n",
       "      <td>A188-1-797667-2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>47</th>\n",
       "      <td>1_148336</td>\n",
       "      <td>A188-1-797667-2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>48</th>\n",
       "      <td>1_148337</td>\n",
       "      <td>A188-1-797667-2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>49</th>\n",
       "      <td>1_148339</td>\n",
       "      <td>A188-1-797667-2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>322</th>\n",
       "      <td>1_1229235</td>\n",
       "      <td>A188-1-797667-3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>324</th>\n",
       "      <td>1_1229248</td>\n",
       "      <td>A188-1-797667-3</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             0                1\n",
       "34    1_147027  A188-1-797667-1\n",
       "36    1_147057  A188-1-797667-1\n",
       "37    1_147071  A188-1-797667-1\n",
       "39    1_147088  A188-1-797667-1\n",
       "46    1_148314  A188-1-797667-2\n",
       "47    1_148336  A188-1-797667-2\n",
       "48    1_148337  A188-1-797667-2\n",
       "49    1_148339  A188-1-797667-2\n",
       "322  1_1229235  A188-1-797667-3\n",
       "324  1_1229248  A188-1-797667-3"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_snp_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "231fe543-f255-4b24-86ea-5d68b963a1a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_snp_df.set_index(0, inplace=True)\n",
    "all_snp_df[\"Tag\"] = \"N\"\n",
    "\n",
    "def parse_tag_file():\n",
    "    skiprows = 0\n",
    "    file_handle = open(tag_file, \"r\")\n",
    "\n",
    "    while True:\n",
    "        line = file_handle.readline()\n",
    "        if \"Test\tAlleles Captured\" in line:\n",
    "            break\n",
    "        else:\n",
    "            skiprows+=1\n",
    "\n",
    "    tag_df = pd.read_csv(tag_file, sep=\"\\t\", skiprows=skiprows)\n",
    "\n",
    "    for index, row in tag_df.iterrows():\n",
    "        if len(row[\"Alleles Captured\"].split(\",\")) > 1:\n",
    "            all_snp_df.at[row[\"Test\"], 'Tag'] = \"Y\"\n",
    "    pass\n",
    "\n",
    "parse_tag_file()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "id": "1f043338-7dd0-4b20-8794-b246c94bc649",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def parse_block(sv_id):\n",
    "    \"\"\"\n",
    "    解析haploview的结果\n",
    "    \"\"\"\n",
    "    block_file = f\"./files/{sv_id}/haploview.GABRIELblocks\"\n",
    "    info_file = f\"./files/{sv_id}/SV.plink.info\"\n",
    "    all_snp_df = pd.DataFrame()\n",
    "    all_block_context = []\n",
    "    \n",
    "    if os.path.exists(block_file) and os.path.exists(info_file):\n",
    "        info_df = pd.read_csv(info_file, header=None, sep=\"\\t\")\n",
    "        block_count = 0\n",
    "        with open(block_file, \"r\") as f:\n",
    "            lines = f.readlines()\n",
    "            block_context = []\n",
    "            block_id = \"\"\n",
    "            for line in lines:\n",
    "                line_text = line.strip()\n",
    "                if \"BLOCK\" in line_text:\n",
    "                    block_count+=1\n",
    "                    block_id = f\"{sv_id}-{block_count}\"\n",
    "                    line_list = line_text.split(\" \")\n",
    "                    marker_index = line_list.index(\"MARKERS:\")\n",
    "                    snp_indexs = [int(v)-1 for v in line_list[marker_index+1:]]\n",
    "                    item_snp_df = info_df.iloc[snp_indexs][[0]].copy()\n",
    "                    item_snp_df[1] = block_id\n",
    "    \n",
    "                    if len(all_snp_df) > 0:\n",
    "                        all_snp_df = pd.concat([all_snp_df, item_snp_df])\n",
    "                    else:\n",
    "                        all_snp_df = item_snp_df\n",
    "    \n",
    "                    if len(block_context) > 0:\n",
    "                        all_block_context += block_context\n",
    "                        block_context = []\n",
    "                else:\n",
    "                    if len(block_context) == 0 and block_id !=\"\":\n",
    "                        block_context.append(f\">{block_id}\")\n",
    "                        block_context.append(line_text)\n",
    "                    else:\n",
    "                        block_context.append(line_text)\n",
    "            if len(block_context) > 0:\n",
    "                all_block_context += block_context\n",
    "    \n",
    "    else:\n",
    "        clear_haploview(sv_id)\n",
    "\n",
    "    return all_snp_df, all_block_context\n",
    "\n",
    "\n",
    "def find_tag_snp(sv_id, sv_snp_df):\n",
    "    tag_file = f\"./files/{sv_id}/haploview.TAGS\"\n",
    "    snp_tag_df = sv_snp_df.set_index(0)\n",
    "    snp_tag_df[\"Tag\"] = \"N\"\n",
    "    \n",
    "    skiprows = 0\n",
    "    target_line_exists = False\n",
    "    \n",
    "    with open(tag_file, \"r\") as f:\n",
    "        for line in f:\n",
    "            if \"Test\tAlleles Captured\" in line:\n",
    "                target_line_exists = True\n",
    "                break\n",
    "            else:\n",
    "                skiprows+=1\n",
    "\n",
    "    if target_line_exists:\n",
    "        tag_df = pd.read_csv(tag_file, sep=\"\\t\", skiprows=skiprows)\n",
    "        for index, row in tag_df.iterrows():\n",
    "            if len(row[\"Alleles Captured\"].split(\",\")) > 1:\n",
    "                snp_tag_df.at[row[\"Test\"], 'Tag'] = \"Y\"\n",
    "    \n",
    "    snp_tag_df.reset_index(inplace=True)\n",
    "\n",
    "    return snp_tag_df\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 133,
   "id": "946d3714-5897-4662-ab41-7a2d574c4bf2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "             0            1\n",
      "282  1_8624633  1_7676397-1\n",
      "284  1_8624644  1_7676397-1\n",
      "287  1_8624670  1_7676397-1\n",
      "289  1_8624684  1_7676397-2\n",
      "290  1_8624688  1_7676397-2 ['>1_7676397-1', '222 (0.869)\\t|0.798\\t0.069|', '444 (0.093)\\t|0.082\\t0.011|', '442 (0.020)\\t|0.017\\t0.003|', '242 (0.017)\\t|0.017\\t0.000|', 'Multiallelic Dprime: 0.064', '>1_7676397-2', '43 (0.915)', '14 (0.083)'] \n"
     ]
    }
   ],
   "source": [
    "\n",
    "sv_id = \"1_7676397\"\n",
    "\n",
    "def parse_block(sv_id, sv_start, sv_stop):\n",
    "    \"\"\"\n",
    "    解析haploview的结果\n",
    "    \"\"\"\n",
    "    block_file = f\"./files/{sv_id}/haploview.GABRIELblocks\"\n",
    "    info_file = f\"./files/{sv_id}/SV.plink.info\"\n",
    "    all_snp_df = pd.DataFrame()\n",
    "    all_block_context = []\n",
    "\n",
    "    include_sv_block=\"\"\n",
    "\n",
    "    if os.path.exists(block_file) and os.path.exists(info_file):\n",
    "        info_df = pd.read_csv(info_file, header=None, sep=\"\\t\")\n",
    "        block_count = 0\n",
    "        with open(block_file, \"r\") as f:\n",
    "            lines = f.readlines()\n",
    "            block_context = []\n",
    "            block_id = \"\"\n",
    "            for line in lines:\n",
    "                line_text = line.strip()\n",
    "                if \"BLOCK\" in line_text:\n",
    "                    block_count+=1\n",
    "                    block_id = f\"{sv_id}-{block_count}\"\n",
    "                    line_list = line_text.split(\" \")\n",
    "                    marker_index = line_list.index(\"MARKERS:\")\n",
    "                    snp_indexs = [int(v)-1 for v in line_list[marker_index+1:]]\n",
    "                    item_snp_df = info_df.iloc[snp_indexs][[0]].copy()\n",
    "                    item_snp_df[1] = block_id\n",
    "\n",
    "                    # check if sv in block\n",
    "                    block_range = [int(info_df.iloc[snp_indexs[0]][1]), int(info_df.iloc[snp_indexs[-1]][1])]\n",
    "                    if block_range[0] <= sv_start and block_range[1] >= sv_stop:\n",
    "                        include_sv_block = block_id\n",
    "\n",
    "                    if len(all_snp_df) > 0:\n",
    "                        all_snp_df = pd.concat([all_snp_df, item_snp_df])\n",
    "                    else:\n",
    "                        all_snp_df = item_snp_df\n",
    "\n",
    "                    if len(block_context) > 0:\n",
    "                        all_block_context += block_context\n",
    "                        block_context = []\n",
    "                else:\n",
    "                    if len(block_context) == 0 and block_id !=\"\":\n",
    "                        block_context.append(f\">{block_id}\")\n",
    "                        block_context.append(line_text)\n",
    "                    else:\n",
    "                        block_context.append(line_text)\n",
    "            if len(block_context) > 0:\n",
    "                all_block_context += block_context\n",
    "\n",
    "    return all_snp_df, all_block_context, include_sv_block\n",
    "\n",
    "\n",
    "all_snp_df, all_block_context, include_sv_block = parse_block(sv_id, 2370620, 2370621)\n",
    "\n",
    "print(all_snp_df, all_block_context, include_sv_block)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "id": "40d90bcb-96b6-4250-af19-7e4d20b47cc0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>0</th>\n",
       "      <th>1</th>\n",
       "      <th>Tag</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1_8624633</td>\n",
       "      <td>1_7676397-1</td>\n",
       "      <td>Y</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1_8624644</td>\n",
       "      <td>1_7676397-1</td>\n",
       "      <td>N</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1_8624670</td>\n",
       "      <td>1_7676397-1</td>\n",
       "      <td>N</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1_8624684</td>\n",
       "      <td>1_7676397-2</td>\n",
       "      <td>N</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1_8624688</td>\n",
       "      <td>1_7676397-2</td>\n",
       "      <td>Y</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           0            1 Tag\n",
       "0  1_8624633  1_7676397-1   Y\n",
       "1  1_8624644  1_7676397-1   N\n",
       "2  1_8624670  1_7676397-1   N\n",
       "3  1_8624684  1_7676397-2   N\n",
       "4  1_8624688  1_7676397-2   Y"
      ]
     },
     "execution_count": 144,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def find_tag_snp(sv_id, sv_snp_df):\n",
    "    \"\"\"\n",
    "    找出block中的Tag SNP\n",
    "    \"\"\"\n",
    "    tag_file = f\"./files/{sv_id}/haploview.TAGS\"\n",
    "    snp_tag_df = sv_snp_df.set_index(0)\n",
    "    snp_tag_df[\"Tag\"] = \"N\"\n",
    "\n",
    "    skiprows = 0\n",
    "    target_line_exists = False\n",
    "\n",
    "    with open(tag_file, \"r\") as f:\n",
    "        for line in f:\n",
    "            if \"Test\tAlleles Captured\" in line:\n",
    "                target_line_exists = True\n",
    "                break\n",
    "            else:\n",
    "                skiprows+=1\n",
    "\n",
    "    # check if target line exists\n",
    "    if target_line_exists:\n",
    "        tag_df = pd.read_csv(tag_file, sep=\"\\t\", skiprows=skiprows)\n",
    "        for index, row in tag_df.iterrows():\n",
    "            if row[\"Test\"] in snp_tag_df.index:\n",
    "                if len(row[\"Alleles Captured\"].split(\",\")) > 1:\n",
    "                    snp_tag_df.at[row[\"Test\"], 'Tag'] = \"Y\"\n",
    "\n",
    "    snp_tag_df.reset_index(inplace=True)\n",
    "    return snp_tag_df\n",
    "\n",
    "\n",
    "tag_df = find_tag_snp(sv_id, all_snp_df)\n",
    "tag_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6552c9c2-6cf3-4817-b958-649c985495bb",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
